import os
import sys
sys.path.append('/home/tfj/code/py/cv/image_retrieval/')
import pandas as pd
import torch
from torch.nn.functional import normalize
from utils.parquet import read_embeddings_from_parquet,save_embeddings_to_parquet

files = [
    'outputs/embeddings/no_train/baseline_embeddings_o.parquet',
    'outputs/embeddings/no_train/baseline_embeddings_q.parquet',
    'outputs/embeddings/no_train/baseline_embeddings_v.parquet',
]

# 读取所有embedding
o = read_embeddings_from_parquet(files[0])
# q = read_embeddings_from_parquet(files[1])
# v = read_embeddings_from_parquet(files[2])

def normalize_embeddings(data_dict):
    for key, item in data_dict.items():
        embedding = item['embedding']
        if not isinstance(embedding, torch.Tensor):
            embedding = torch.tensor(embedding, dtype=torch.float32)
        embedding = normalize(embedding.unsqueeze(0), p=2, dim=1)  # (1, D)
        data_dict[key]['embedding'] = embedding.squeeze(0).cpu().numpy()  # 回到 (D,)
        data_dict[key]['image_id'] = key
    return data_dict
# 对每个集合进行归一化

new_o = normalize_embeddings(o)
# new_q = normalize_embeddings(q)
# new_v = normalize_embeddings(v)

save_embeddings_to_parquet(new_o,files[0].replace('.parquet','_norm.parquet'))
# save_embeddings_to_parquet(new_q,files[1].replace('.parquet','_norm.parquet'))
# save_embeddings_to_parquet(new_v,files[2].replace('.parquet','_norm.parquet'))




print(1)
